import sys
import os
import numpy as np
from PIL import Image  # 添加PIL的Image导入
from PyQt5.QtWidgets import (QApplication, QMainWindow, QLabel, QAction, QFileDialog, 
                            QMessageBox, QMenu, QVBoxLayout, QWidget, QPushButton)
from PyQt5.QtGui import QPixmap, QImage, QPainter, QPen, QColor
from PyQt5.QtCore import Qt, QRect

import image_io
import image_transform
import image_channels
from image_utils import logging, language_data
import config

class ImageProcessorQt(QMainWindow):
    def __init__(self):
        super().__init__()
        
        # 初始化变量
        self.image = None
        self.image_path = None
        self.current_language = "zh"  # 默认语言为中文
        self.is_drawing = False
        self.current_rect = None
        self.rect_start_x = 0
        self.rect_start_y = 0
        self.rect_end_x = 0
        self.rect_end_y = 0
        self.crop_region = None
        self.temp_rect = None
        
        self.init_ui()
        
    def init_ui(self):
        """初始化UI"""
        self.setWindowTitle("图像处理应用 (Qt版)")
        self.setGeometry(100, 100, 800, 600)
        
        # 创建中央窗口部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        # 创建布局
        layout = QVBoxLayout(central_widget)
        
        # 创建图像显示标签
        self.image_label = QLabel()
        self.image_label.setAlignment(Qt.AlignCenter)
        self.image_label.setMinimumSize(600, 400)
        self.image_label.setStyleSheet("border: 1px solid #cccccc;")
        self.image_label.mousePressEvent = self.on_mouse_down
        self.image_label.mouseMoveEvent = self.on_mouse_move
        self.image_label.mouseReleaseEvent = self.on_mouse_up
        
        layout.addWidget(self.image_label)
        
        # 创建菜单栏
        self.create_menus()
        
        self.show()
        
    def create_menus(self):
        """创建菜单栏"""
        # 文件菜单
        file_menu = self.menuBar().addMenu(language_data[self.current_language]['file'])
        
        open_original_action = QAction(language_data[self.current_language]['open_original'], self)
        open_original_action.triggered.connect(self.open_as_original)
        file_menu.addAction(open_original_action)
        
        open_gray_action = QAction(language_data[self.current_language]['open_grayscale'], self)
        open_gray_action.triggered.connect(self.open_as_gray)
        file_menu.addAction(open_gray_action)
        
        save_action = QAction(language_data[self.current_language]['save_image'], self)
        save_action.triggered.connect(self.save_image)
        file_menu.addAction(save_action)
        
        # 编辑菜单
        edit_menu = self.menuBar().addMenu(language_data[self.current_language]['edit'])
        
        crop_action = QAction(language_data[self.current_language]['crop'], self)
        crop_action.triggered.connect(self.crop_image)
        edit_menu.addAction(crop_action)
        
        reset_action = QAction(language_data[self.current_language]['reset'], self)
        reset_action.triggered.connect(self.reset_image)
        edit_menu.addAction(reset_action)
        
        # 信息菜单
        info_menu = self.menuBar().addMenu(language_data[self.current_language]['info'])
        
        details_action = QAction(language_data[self.current_language]['details'], self)
        details_action.triggered.connect(self.show_image_details)
        info_menu.addAction(details_action)
        
        # 通道菜单
        channel_menu = self.menuBar().addMenu(language_data[self.current_language]['split'])
        
        split_action = QAction(language_data[self.current_language]['split_channels'], self)
        split_action.triggered.connect(self.split_image_channels)
        channel_menu.addAction(split_action)
        
        merge_action = QAction(language_data[self.current_language]['merge_channels'], self)
        merge_action.triggered.connect(self.merge_image_channels)
        channel_menu.addAction(merge_action)
        
    def open_as_original(self):
        """打开原始图像"""
        logging.info("正在打开原始图像...")
        file_path, _ = QFileDialog.getOpenFileName(self, "打开图像", "", "图像文件 (*.png *.jpg *.jpeg *.bmp *.tif)")
        
        if not file_path:
            logging.info("没有选择任何图像文件,打开操作已取消。")
            return
            
        self.image_path = file_path
        self.image = image_io.load_image(self.image_path)
        self.display_image()
        logging.info("成功打开原始图像。")
        
    def open_as_gray(self):
        """将原始图像转换为灰度图像"""
        file_path, _ = QFileDialog.getOpenFileName(self, "打开图像", "", "图像文件 (*.png *.jpg *.jpeg *.bmp *.tif)")
        
        if not file_path:
            logging.info("没有选择任何图像文件,打开操作已取消。")
            return
            
        self.image_path = file_path
        self.image = image_io.load_image(self.image_path, as_gray=True)
        self.display_image()
        logging.info("成功打开并转换为灰度图像。")
        
    def display_image(self):
        """在标签上显示当前图像"""
        if self.image is None:
            logging.info("当前没有图像可显示。")
            return
            
        # 将numpy数组转换为QImage
        height, width = self.image.shape[:2]
        
        if len(self.image.shape) == 2:  # 灰度图像
            bytes_per_line = width
            q_image = QImage(self.image.data, width, height, bytes_per_line, QImage.Format_Grayscale8)
        else:  # 彩色图像
            bytes_per_line = 3 * width
            # 注意：将numpy数组转换为PIL图像，然后转换为QImage
            pil_img = Image.fromarray(self.image)
            q_image = QImage(pil_img.tobytes(), width, height, bytes_per_line, QImage.Format_RGB888)
            
        # 创建QPixmap并调整大小以适应标签
        pixmap = QPixmap.fromImage(q_image)
        
        # 计算缩放比例
        label_width = self.image_label.width()
        label_height = self.image_label.height()
        
        if width > label_width or height > label_height:
            pixmap = pixmap.scaled(label_width, label_height, Qt.KeepAspectRatio, Qt.SmoothTransformation)
            
        self.image_label.setPixmap(pixmap)
        logging.info("成功显示图像。")
        
    def show_image_details(self):
        """显示当前图像的详细信息"""
        if self.image is None:
            QMessageBox.information(self, language_data[self.current_language]['error'], 
                                  language_data[self.current_language]['no_image_loaded'])
            logging.info("没有载入任何图像,无法显示图像细节。")
            return
            
        # 从路径中提取文件名
        filename = os.path.basename(self.image_path)
        
        # 计算其他基本信息
        total_pixels = self.image.size  # 获取图像的总像素数
        shape = self.image.shape  # 获取图像的形状(高度,宽度,通道数)
        channels = 1 if len(self.image.shape) == 2 else self.image.shape[2]  # 获取图像的通道数
        dtype = self.image.dtype  # 获取图像像素的数据类型
        
        # 构建信息字符串,包括文件名
        details = language_data[self.current_language]['info_str'].format(filename, total_pixels, shape, channels, dtype)
        
        # 显示信息
        QMessageBox.information(self, "INFO", details)
        logging.info("已显示图像细节。")
        
    def on_mouse_down(self, event):
        """处理鼠标左键按下事件"""
        if self.image is None:
            return
            
        self.is_drawing = True
        self.rect_start_x = event.x()
        self.rect_start_y = event.y()
        self.temp_rect = None
        
    def on_mouse_move(self, event):
        """处理鼠标左键移动事件"""
        if not self.is_drawing:
            return
            
        # 更新临时矩形
        self.rect_end_x = event.x()
        self.rect_end_y = event.y()
        
        # 强制重绘
        self.update()
        
    def on_mouse_up(self, event):
        """处理鼠标左键释放事件"""
        if not self.is_drawing:
            return
            
        self.is_drawing = False
        self.rect_end_x = event.x()
        self.rect_end_y = event.y()
        
        # 计算图像上的实际裁剪区域
        pixmap = self.image_label.pixmap()
        if pixmap:
            # 获取图像在标签中的位置和大小
            img_rect = pixmap.rect()
            img_x = (self.image_label.width() - img_rect.width()) // 2
            img_y = (self.image_label.height() - img_rect.height()) // 2
            
            # 计算相对于图像的坐标
            rel_start_x = max(0, self.rect_start_x - img_x)
            rel_start_y = max(0, self.rect_start_y - img_y)
            rel_end_x = min(img_rect.width(), self.rect_end_x - img_x)
            rel_end_y = min(img_rect.height(), self.rect_end_y - img_y)
            
            # 计算缩放比例
            scale_x = self.image.shape[1] / img_rect.width()
            scale_y = self.image.shape[0] / img_rect.height()
            
            # 计算实际图像上的裁剪区域
            img_start_x = int(rel_start_x * scale_x)
            img_start_y = int(rel_start_y * scale_y)
            img_end_x = int(rel_end_x * scale_x)
            img_end_y = int(rel_end_y * scale_y)
            
            # 保存裁剪区域
            self.crop_region = (
                min(img_start_y, img_end_y),
                max(img_start_y, img_end_y),
                min(img_start_x, img_end_x),
                max(img_start_x, img_end_x)
            )
            
        # 强制重绘
        self.update()
        logging.info("已绘制裁剪区域。")
        
    def paintEvent(self, event):
        """重写绘制事件以显示裁剪框"""
        super().paintEvent(event)
        
        if self.is_drawing and self.image is not None:
            painter = QPainter(self)
            pen = QPen(QColor(255, 0, 0))  # 红色
            pen.setWidth(2)
            painter.setPen(pen)
            
            # 绘制矩形
            rect = QRect(
                min(self.rect_start_x, self.rect_end_x),
                min(self.rect_start_y, self.rect_end_y),
                abs(self.rect_end_x - self.rect_start_x),
                abs(self.rect_end_y - self.rect_start_y)
            )
            painter.drawRect(rect)
            
    def save_image(self):
        """保存当前图像"""
        if self.image is None:
            QMessageBox.information(self, language_data[self.current_language]['error'], 
                                  language_data[self.current_language]['no_image_to_be_saved'])
            logging.error("没有图像可保存,保存操作已取消。")
            return
            
        file_path, _ = QFileDialog.getSaveFileName(self, "保存图像", "", 
                                                "PNG (*.png);;JPEG (*.jpg);;BMP (*.bmp);;TIFF (*.tif)")
        
        if not file_path:
            logging.info("没有选择保存路径,保存操作已取消。")
            return
            
        image_io.save_image(self.image, file_path)
        logging.info(f"图像已成功保存到 {file_path}。")
        
    def crop_image(self):
        """裁剪当前图像"""
        if self.image is None:
            QMessageBox.information(self, language_data[self.current_language]['error'], 
                                  language_data[self.current_language]['no_image_to_be_cropped'])
            logging.error("没有图像可裁剪,裁剪操作已取消。")
            return
            
        # 检查是否有绘制的裁剪区域
        if hasattr(self, 'crop_region') and self.crop_region:
            cropped_image = image_transform.crop_image(self.image, self.crop_region)
            
            if cropped_image is not None:
                self.image = cropped_image
                self.display_image()
                logging.info("图像裁剪成功。")
                
                # 清零裁剪框信息
                self.crop_region = None
                self.update()
            else:
                QMessageBox.information(self, language_data[self.current_language]['error'], 
                                      language_data[self.current_language]['crop_failed'])
                logging.error("裁剪操作失败。")
        else:
            QMessageBox.information(self, language_data[self.current_language]['error'], 
                                  language_data[self.current_language]['no_crop_area'])
            logging.error("没有选择裁剪区域,裁剪操作已取消。")
            
    def reset_image(self):
        """重置当前图像为原始状态"""
        if not self.image_path:
            QMessageBox.information(self, language_data[self.current_language]['error'], 
                                  language_data[self.current_language]['no_reset_image'])
            logging.error("没有图像可重置,重置操作已取消。")
            return
            
        self.image = image_io.load_image(self.image_path)
        self.display_image()
        logging.info("图像已重置为原始状态。")
        
    def split_image_channels(self):
        """分割当前图像的通道,并在新窗口中显示每个通道"""
        if self.image is None:
            QMessageBox.information(self, language_data[self.current_language]['error'], 
                                  language_data[self.current_language]['no_image_loaded'])
            return
            
        if len(self.image.shape) == 2:
            QMessageBox.information(self, language_data[self.current_language]['error'], 
                                  language_data[self.current_language]['single_channel_error'])
            return
        
        try:
            # 分割通道
            channels = image_channels.split_image_channels(self.image)
            
            # 创建并保存对象引用，防止被垃圾回收
            self.channel_windows = []
            self.channel_pixmaps = []
            self.channel_arrays = channels
            
            # 为每个通道创建一个新窗口
            for i, channel in enumerate(channels):
                # 创建子窗口
                channel_window = QWidget()
                channel_window.setWindowTitle(f"通道 {i}")
                channel_window.setGeometry(100 + i*50, 100 + i*50, 400, 400)
                
                # 创建布局
                layout = QVBoxLayout()
                channel_window.setLayout(layout)
                
                # 添加标题
                title_label = QLabel(f"通道 {i}")
                title_label.setAlignment(Qt.AlignCenter)
                title_label.setStyleSheet("font-size: 15px; font-weight: bold;")
                layout.addWidget(title_label)
                
                # 创建图像标签
                channel_label = QLabel()
                channel_label.setAlignment(Qt.AlignCenter)
                channel_label.setMinimumSize(300, 300)
                layout.addWidget(channel_label)
                
                # 创建保存按钮
                save_button = QPushButton("保存子图")
                save_button.setFixedWidth(200)
                
                # 创建一个局部函数来保存单个通道
                def create_save_function(idx):
                    def save_channel():
                        save_path, _ = QFileDialog.getSaveFileName(
                            channel_window, 
                            "保存通道图像", 
                            "", 
                            "PNG (*.png);;JPEG (*.jpg);;BMP (*.bmp)"
                        )
                        if save_path:
                            image_io.save_image(self.channel_arrays[idx], save_path)
                            QMessageBox.information(
                                channel_window, 
                                "成功", 
                                f"通道 {idx} 已保存为 {save_path}"
                            )
                    return save_channel
                
                save_button.clicked.connect(create_save_function(i))
                layout.addWidget(save_button, alignment=Qt.AlignCenter)
                
                # 显示通道图像
                # 注意：将numpy数组转换为8位无符号整数类型
                channel_8bit = channel
                if channel.dtype != np.uint8:
                    if channel.max() <= 1.0:
                        channel_8bit = (channel * 255).astype(np.uint8)
                    else:
                        channel_8bit = channel.astype(np.uint8)
                
                # 使用PIL创建图像
                pil_img = Image.fromarray(channel_8bit)
                qimg = QImage(pil_img.tobytes(), pil_img.width, pil_img.height, pil_img.width, QImage.Format_Grayscale8)
                pixmap = QPixmap.fromImage(qimg)
                
                # 缩放图像以适应标签
                pixmap = pixmap.scaled(300, 300, Qt.KeepAspectRatio, Qt.SmoothTransformation)
                channel_label.setPixmap(pixmap)
                
                # 将窗口和图像添加到列表中，防止被垃圾回收
                self.channel_windows.append(channel_window)
                self.channel_pixmaps.append(pixmap)
                
                # 显示窗口
                channel_window.show()
            
            logging.info(f"成功分割{len(channels)}个通道")
            
        except Exception as e:
            QMessageBox.critical(self, "错误", str(e))
            logging.error(f"分离通道时出错: {e}")
            
    def merge_image_channels(self):
        """合并多个单通道图像为一个多通道图像"""
        # 打开文件对话框,让用户选择多个图像文件
        file_paths, _ = QFileDialog.getOpenFileNames(self, "选择要合并的通道图像", "", 
                                                 "图像文件 (*.png *.jpg *.jpeg *.bmp *.tif)")
        
        # 如果用户没有选择文件,直接返回
        if not file_paths:
            logging.info("没有选择任何通道图像,合并操作已取消。")
            return
            
        # 如果只选择了一个文件,显示错误信息并返回
        if len(file_paths) == 1:
            QMessageBox.critical(self, language_data[self.current_language]['error'], 
                               language_data[self.current_language]['single_file_error'])
            return
            
        try:
            # 合并通道图像
            self.image = image_channels.merge_image_channels(file_paths)
            
            # 显示合并后的图像
            self.display_image()
            
            # 更新图像路径（但不保存，除非用户选择保存）
            self.image_path = ""
            
            logging.info("成功合并通道")
            
        except Exception as e:
            QMessageBox.critical(self, "错误", str(e))
            logging.error(f"合并通道时出错: {e}")

def main():
    app = QApplication(sys.argv)
    window = ImageProcessorQt()
    sys.exit(app.exec_())

if __name__ == "__main__":
    main()
